import asyncio
import threading

import arrow
import redis
import redis.asyncio as async_redis
from motor.core import AgnosticCollection
from motor.motor_asyncio import AsyncIOMotorClient

from fastapi_template.config import app_config


def get_thread_context():
    _thread_context = threading.local()
    if not hasattr(_thread_context, "MONGODB_ASYNC_CLIENTS"):
        _thread_context.MONGODB_ASYNC_CLIENTS = dict()
    if not hasattr(_thread_context, "REDIS_ASYNC_CLIENTS"):
        _thread_context.REDIS_ASYNC_CLIENTS = dict()
    if not hasattr(_thread_context, "REDIS_CLIENTS"):
        _thread_context.REDIS_CLIENTS = dict()
    return _thread_context


def get_async_mongo_client(client_name="default") -> AsyncIOMotorClient:
    """
    :return: MongoClient实例
    """
    _thread_context = get_thread_context()
    if client_name not in _thread_context.MONGODB_ASYNC_CLIENTS:
        conf_key = "%s_MONGO_URI" % client_name.upper()
        uri = str(getattr(app_config, conf_key))
        connect_params = dict()
        connect_params["connectTimeoutMS"] = 5000
        connect_params["serverSelectionTimeoutMS"] = 7000
        client = AsyncIOMotorClient(uri, **connect_params)
        _thread_context.MONGODB_ASYNC_CLIENTS[client_name] = client

    return _thread_context.MONGODB_ASYNC_CLIENTS[client_name]


def get_async_mongo_db(db=None, client_name="default"):
    """
    获取 Mongo DB
    :return: MongoDatabase
    """
    client = get_async_mongo_client(client_name)
    conf_key = "%s_MONGO_DBNAME" % client_name.upper()
    db = db or getattr(app_config, conf_key)
    return client[db]


def get_async_mongo_collection(collection,
                               db=None,
                               client_name="default") -> AgnosticCollection:
    """
    获取 Mongo Collection
    :return: MongoCollection
    """
    client = get_async_mongo_client(client_name)
    """
    local_timezone = app_config.get('LOCAL_TIMEZONE')
    codec_options = CodecOptions(
        tz_aware=True, tzinfo=pytz.timezone(local_timezone)
    )
    """
    conf_key = "%s_MONGO_DBNAME" % client_name.upper()
    db = db or getattr(app_config, conf_key)
    return client[db][collection].with_options()


def get_redis_client(client_name="default"):
    """
    :return: RedisClient实例
    """
    _thread_context = get_thread_context()
    db = 8 if app_config.ENV == "prod" else 0
    key = "%s_%s" % (client_name, db)
    if key not in _thread_context.REDIS_CLIENTS:
        host = app_config.REDIS_DSN.host
        port = app_config.REDIS_DSN.port or 6379
        _thread_context.REDIS_CLIENTS[key] = redis.Redis(host=host,
                                                         port=port,
                                                         db=db,
                                                         protocol=3)
    return _thread_context.REDIS_CLIENTS[key]


def get_async_redis_client(client_name="default") -> async_redis.Redis:
    """
    :return: RedisClient实例
    """
    _thread_context = get_thread_context()
    db = app_config.CACHE_DB
    key = "async_%s" % client_name
    if key not in _thread_context.REDIS_ASYNC_CLIENTS:
        host = app_config.REDIS_DSN.host
        port = app_config.REDIS_DSN.port or 6379
        _thread_context.REDIS_ASYNC_CLIENTS[key] = async_redis.Redis(
            host=host, port=port, db=db, protocol=3)
    return _thread_context.REDIS_ASYNC_CLIENTS[key]


def get_cache_key(tpl: str, *args) -> str:
    """simple cache key generator
    :param tpl: cache key template
    :param args: cache key template args
    :return: cache key

    Example:
    >>> get_cache_key("user:%s", "user_id")
    "HAS:user:user_id"
    """
    key = tpl % args
    return "%s:%s" % (app_config.CACHE_NAMESPACE, key)


async def _test_connection():
    collection2 = get_async_mongo_collection("fastapi_test_temp")
    print(collection2.name)

    documents = []
    now = arrow.now()
    for _ in range(100):
        now = now.shift(minutes=-_)
        document = {"key": "value", "create_time": now.datetime}
        documents.append(document)
    result = await collection2.insert_many(documents)
    print(result)

    # redis_client = get_async_redis_client()
    # key = "fastapi-test-temp-key-00001"
    # await redis_client.set(name=key, value=b'1', ex=10)
    # got = await redis_client.get(name=key)
    # assert got == b'1'


if __name__ == "__main__":
    asyncio.run(_test_connection())
